import asyncio
import json
from typing import Dict
from datetime import datetime
from .queue_manager import MessageQueueManager
from .state_manager import state_manager  # 直接导入实例

class AsyncSocketManager:
    def __init__(self, host: str = '0.0.0.0', port: int = 9999):
        self.host = host
        self.port = port
        self.server = None
        self.clients: Dict[str, asyncio.StreamWriter] = {}
        self.running = False
        self.queue_manager = MessageQueueManager()
        # 在状态管理器中记录服务器状态
        state_manager.set("socket_server", {
            "host": host,
            "port": port,
            "status": "initialized"
        })
    
    def json_dumps(self, obj) -> str:
        """JSON编码，支持中文"""
        return json.dumps(obj, ensure_ascii=False)
    
    async def start(self):
        """启动服务器"""
        try:
            # 更新服务器状态
            state_manager.set("socket_server", {
                "status": "running",
                "started_at": datetime.now().isoformat()
            })
            self.server = await asyncio.start_server(
                self.handle_client,
                self.host,
                self.port
            )
            self.running = True
            print(f"Socket服务器启动在 {self.host}:{self.port}")
            
            async with self.server:
                await self.server.serve_forever()
                
        except Exception as e:
            print(f"服务器启动失败: {str(e)}")
            self.running = False
    
    async def stop(self):
        """停止服务器"""
        self.running = False
        if self.server:
            self.server.close()
            await self.server.wait_closed()
            
        # 关闭所有客户端连接
        for writer in self.clients.values():
            writer.close()
            await writer.wait_closed()
        print("服务器已停止")
    
    async def handle_client(self, reader: asyncio.StreamReader, writer: asyncio.StreamWriter):
        """处理客户端连接和消息"""
        peer = writer.get_extra_info('peername')
        client_id = f"{peer[0]}:{peer[1]}"
        
        try:
            # 存储客户端连接
            self.clients[client_id] = writer
            print(f"新客户端连接: {client_id}")
            
            # 发送欢迎消息
            welcome_msg = self.json_dumps({
                "type": "system",
                "message": "连接成功！",
                "timestamp": datetime.now().isoformat()
            })
            writer.write(welcome_msg.encode('utf-8'))
            await writer.drain()
            
            # 订阅消息队列
            subscriber_queue = await self.queue_manager.subscribe(client_id)
            
            while self.running:
                try:
                    # 接收消息
                    data = await reader.read(4096)
                    if not data:
                        break
                    
                    message = data.decode('utf-8')
                    print(f"收到来自 {client_id} 的消息: {message}")
                    
                    # 发布到消息队列
                    await self.queue_manager.publish(message, f"socket-{client_id}")
                    
                    try:
                        # 尝试解析JSON消息
                        message_data = json.loads(message)
                        message_type = message_data.get("type", "message")
                        
                        if message_type == "heartbeat":
                            # 处理心跳消息
                            response = self.json_dumps({
                                "type": "heartbeat",
                                "timestamp": datetime.now().isoformat()
                            })
                            writer.write(response.encode('utf-8'))
                            await writer.drain()
                        else:
                            # 处理普通消息
                            # 发送确认消息
                            ack_message = self.json_dumps({
                                "type": "ack",
                                "message": "消息已收到",
                                "original_message": message_data.get("message", ""),
                                "timestamp": datetime.now().isoformat()
                            })
                            writer.write(ack_message.encode('utf-8'))
                            await writer.drain()
                            
                            # 广播消息给其他客户端
                            broadcast_message = self.json_dumps({
                                "type": "message",
                                "client_id": client_id,
                                "message": message_data.get("message", ""),
                                "timestamp": datetime.now().isoformat()
                            })
                            await self.broadcast(broadcast_message, exclude=client_id)
                            
                    except json.JSONDecodeError:
                        # 处理普通文本消息
                        response = self.json_dumps({
                            "type": "message",
                            "client_id": client_id,
                            "message": message,
                            "timestamp": datetime.now().isoformat()
                        })
                        await self.broadcast(response)
                        
                except Exception as e:
                    print(f"处理消息时出错: {str(e)}")
                    break
                    
        finally:
            # 取消订阅
            await self.queue_manager.unsubscribe(client_id)
            await self.disconnect_client(client_id)
    
    async def disconnect_client(self, client_id: str):
        """断开客户端连接"""
        if client_id in self.clients:
            writer = self.clients[client_id]
            writer.close()
            await writer.wait_closed()
            del self.clients[client_id]
            print(f"客户端 {client_id} 已断开连接")
            
            # 广播客户端离开消息
            leave_message = self.json_dumps({
                "type": "system",
                "message": f"Client {client_id} left the chat",
                "timestamp": datetime.now().isoformat()
            })
            await self.broadcast(leave_message)
    
    async def broadcast(self, message: str, exclude: str = None):
        """广播消息给所有客户端"""
        for cid, writer in list(self.clients.items()):
            if cid != exclude:
                try:
                    writer.write(message.encode('utf-8'))
                    await writer.drain()
                except Exception as e:
                    print(f"向客户端 {cid} 发送消息时出错: {str(e)}")
                    await self.disconnect_client(cid) 